
#ifndef __SDK__GENERATOR_BASE_H__INCLUDE__
#define __SDK__GENERATOR_BASE_H__INCLUDE__

#include "BackupManager.h"
#include "ServiceBase.h"
#include "Splitter.h"
#include "Task.h"
#include "TaskProcessor.h"
#include "Utility/Mutex.h"
#include "Utility/TaskScheduler.h"
#include "Utility/TimedAccumulator.h"
#include <map>
#include <set>
#include <string>
#include <vector>

namespace nebula { namespace sdk {

class GeneratorBase : public sdk::ServiceBase
{
public:
	GeneratorBase();
	virtual ~GeneratorBase();

public:
	virtual NEBULA_ERROR_CODE OnInitialize () = 0;
	virtual NEBULA_ERROR_CODE OnStart      () = 0;
	virtual void              OnStop       () = 0;
	virtual void              OnFinalize   () = 0;

	virtual ITask*            AllocateTask () = 0;
	virtual ITask*            ForkTask     (ITask*, const char* id=0) = 0;

protected:
	// called after call to splitter->Split
	virtual void              OnSplit      (ITask*, std::vector<ISplitStrategy*> &strategies);

	// called before a task would been sent/resent to done-queue
	virtual void              OnSendTask   (ITask*, unsigned retries);

	// called before a task would been marked as dead
	virtual void              OnDeadTask   (ITask*);

	// called after an ACK has been received
	virtual void              OnAck        (const char* id, bool effective, int latency);

	// if a task has been marked as dead, Generator will wait for some time and check
	// whether the ACK was lost or timed-out; see EnableDeadTaskTrace for more info
	virtual void              OnAckLost    (const char* id);

	// if a task has been marked as dead, Generator will wait for some time and check
	// whether the ACK was lost or timed-out; see EnableDeadTaskTrace for more info
	virtual void              OnAckTimedOut(const char* id);

	// Overwrite this function to do something when Generator found it is overload
	// and have to pause itself
	virtual void              OnGeneratorPausing ();

	// Overwrite this function to do something when Generator found it is underload
	// and have to resume itself
	virtual void              OnGeneratorResuming();

protected:
	// initialize the generator
	NEBULA_ERROR_CODE InitializeGenerator (const char* backup_dir,
	                                       const char* dead_dir,
	                                       const char* out_slot_name,
	                                       const char* ack_slot_name="ack");

	// set the max-resend count, the default value is 1
	// NOTE: could only be called before StartGenerator
	NEBULA_ERROR_CODE SetMaxResend        (unsigned int);

	// set the resend-timeout, the default value is 120 seconds
	// NOTE: could only be called before StartGenerator
	NEBULA_ERROR_CODE SetResendTimeout    (unsigned int);

	// set the splitter, the default splitter is NULL;
	// GeneratorBase will be responsible to delete the splitter
	// NOTE: could only be called before StartGenerator
	NEBULA_ERROR_CODE SetSplitter         (ISplitter *);

	// enable/disable tracing on dead task; when enabled, Generator will trace the
	// ACK's of the dead tasks, check whether the ACK was lost or just timed-out;
	// The tracing is by default turned off
	NEBULA_ERROR_CODE EnableDeadTaskTrace (bool enable);

	// finalize the generator
	void              FinalizeGenerator   ();

	// start the generator
	NEBULA_ERROR_CODE StartGenerator      ();

	// load legacy tasks from backup dir
	NEBULA_ERROR_CODE LoadLegacyTasks     ();
	NEBULA_ERROR_CODE LoadLegacyTasksInBackground ();

	// stop the generator
	void              StopGenerator       ();

	// after you have receive a task from client, call this function to
	// 1.  Make backup of the task
	// 2.  Do split if any splitter set
	// 3.  Add descendants tasks (or original task, if no splitter set) to out-queue
	NEBULA_ERROR_CODE HandleTask          (ITask *);

	// task allocator adaptor, will call AllocateTask to allocate a new task
	static ITask*     TaskAllocator       (IService*);

	// whether the generator is pausing; see also to OnGeneratorPausing
	// and OnGeneratorResuming
	bool              IsGeneratorPausing  () const;

protected:
	virtual bool      IsOverload          ();
	virtual bool      IsUnderload         ();

private:
	NEBULA_ERROR_CODE ack_handler         (std::auto_ptr<ITask> ack);
	NEBULA_ERROR_CODE do_resends          (bool &repeat_immediately);
	NEBULA_ERROR_CODE split_and_send      (ITask* task);

private:
	bool                                  m_started;
	std::string                           m_backup_dir;
	std::string                           m_dead_dir;
	std::string                           m_out_slot_name;
	std::string                           m_ack_slot_name;
	unsigned int                          m_max_resend;
	unsigned int                          m_resend_timeout;
	std::auto_ptr<ISplitter>              m_splitter;

	std::auto_ptr<IBackupManager>         m_backup_mgr;

	ITodoTaskQueue *                      m_ack_queue;
	IDoneTaskQueue *                      m_out_queue;

	TaskProcessor<GeneratorBase>          m_ack_processor;
	utility::TaskScheduler<GeneratorBase> m_resend_scheduler;

private:
	utility::TaskScheduler<GeneratorBase> m_legacy_tasks_loader;

	NEBULA_ERROR_CODE load_legacy_tasks   (bool &repeat_immediately);
	NEBULA_ERROR_CODE load_legacy_tasks   (int *loaded);

private:
	struct descendant_task_info_t
	{
		std::string     task_id;
		std::string     original_task_id;
		ISplitStrategy  *split_strategy;

		time_t          last_sent;
		unsigned        resend_count;
	};

	typedef std::map<std::string, descendant_task_info_t> descendants_map_t;
	typedef std::map<std::string, std::set<std::string> > split_map_t;

	utility::Mutex                        m_lock;
	descendants_map_t                     m_descendants;
	split_map_t                           m_split_map;

	NEBULA_ERROR_CODE handle_timeout_task (descendant_task_info_t &, bool &is_dead);

private:
	struct dead_task_info_t
	{
		time_t    die_time;
		unsigned  ack_count;
	};

	typedef std::map<std::string, dead_task_info_t> dead_map_t;
	bool                                  m_trace_dead;
	utility::Mutex                        m_dead_lock; 
	dead_map_t                            m_deads;

	utility::TaskScheduler<GeneratorBase> m_trace_dead_scheduler;

	NEBULA_ERROR_CODE trace_dead_tasks    (bool &repeat_immediately);

private:
	bool                                  m_pausing;
	utility::TimedAccumulator             m_effective_ack_statistic;

	void              check_overload      ();
};

}} // namespace nebula::sdk

#endif // __SDK__GENERATOR_BASE_H__INCLUDE__

